#!/bin/bash

# Qwen3 32B 微调快速启动脚本
# 支持自动路径识别和配置驱动

set -e  # 遇到错误立即退出

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 获取脚本所在目录
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"

# 打印带颜色的消息
print_status() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

print_header() {
    echo -e "${BLUE}$1${NC}"
}

# 显示帮助信息
show_help() {
    echo "Qwen3 32B 微调快速启动脚本"
    echo ""
    echo "用法: $0 [选项]"
    echo ""
    echo "选项:"
    echo "  --config              显示当前配置"
    echo "  --check-env           仅检查环境"
    echo "  --download-only       仅下载模型和数据集"
    echo "  --train-only          仅训练（跳过下载）"
    echo "  --test-only           仅测试模型"
    echo "  --quick-test          快速测试（10步训练）"
    echo "  --full-training       完整训练（使用配置中的步数）"
    echo "  --help, -h            显示此帮助信息"
    echo ""
    echo "示例:"
    echo "  $0                    # 完整流程（检查环境 + 下载 + 训练 + 测试）"
    echo "  $0 --quick-test       # 快速测试（10步训练）"
    echo "  $0 --config           # 查看当前配置"
    echo "  $0 --check-env        # 仅检查环境"
    echo ""
}

# 检查环境
check_environment() {
    print_header "🔍 检查环境"
    python "$SCRIPT_DIR/check_environment.py"
    echo ""
}

# 显示配置
show_config() {
    print_header "🔧 当前配置"
    cd "$SCRIPT_DIR"
    python -c "import config; config.print_config()"
    echo ""
}

# 下载模型和数据集
download_resources() {
    print_header "📥 下载资源"
    cd "$SCRIPT_DIR"
    
    print_status "下载模型和数据集..."
    python finetune_qwen3.py --download_model --download_data --download_only
    echo ""
}

# 执行训练
run_training() {
    local max_steps=${1:-""}
    
    print_header "🚀 开始训练"
    cd "$SCRIPT_DIR"
    
    if [[ -n "$max_steps" ]]; then
        print_status "开始训练 (自定义步数: $max_steps)..."
        python finetune_qwen3.py --max_steps "$max_steps" --skip_download
    else
        print_status "开始训练 (使用配置中的步数)..."
        python finetune_qwen3.py --skip_download
    fi
    echo ""
}

# 测试模型
test_model() {
    print_header "🧪 测试模型"
    cd "$SCRIPT_DIR"
    
    print_status "运行测试用例..."
    python test_model.py --mode test
    echo ""
    
    print_status "测试完成! 如需交互对话，请运行:"
    echo "  python $SCRIPT_DIR/test_model.py --mode chat"
    echo ""
}

# 完整流程
full_workflow() {
    local quick_test=${1:-false}
    
    print_header "🎯 Qwen3 32B 微调完整流程"
    echo "项目根目录: $PROJECT_ROOT"
    echo "脚本目录: $SCRIPT_DIR"
    echo ""
    
    # 检查环境
    check_environment
    
    # 显示配置
    show_config
    
    # 下载资源
    download_resources
    
    # 执行训练
    if [[ "$quick_test" == "true" ]]; then
        print_warning "快速测试模式：仅训练10步"
        run_training 10
    else
        run_training
    fi
    
    # 测试模型
    test_model
    
    print_header "🎉 流程完成！"
    print_status "微调已完成，模型已保存到输出目录"
    print_status "可以使用以下命令进行更多测试："
    echo "  cd $SCRIPT_DIR"
    echo "  python test_model.py --mode chat    # 交互对话"
    echo "  python test_model.py --mode benchmark # 性能测试"
}

# 解析命令行参数
case "${1:-}" in
    --help|-h)
        show_help
        exit 0
        ;;
    --config)
        show_config
        exit 0
        ;;
    --check-env)
        check_environment
        exit 0
        ;;
    --download-only)
        download_resources
        exit 0
        ;;
    --train-only)
        run_training
        exit 0
        ;;
    --test-only)
        test_model
        exit 0
        ;;
    --quick-test)
        full_workflow true
        exit 0
        ;;
    --full-training)
        full_workflow false
        exit 0
        ;;
    "")
        # 默认执行完整流程
        full_workflow false
        exit 0
        ;;
    *)
        print_error "未知选项: $1"
        echo ""
        show_help
        exit 1
        ;;
esac
import torch
import transformers
import datasets
import peft
import bitsandbytes
import modelscope
print('✓ 所有必要的库都已安装')
print(f'✓ PyTorch版本: {torch.__version__}')
print(f'✓ Transformers版本: {transformers.__version__}')
print(f'✓ ModelScope版本: {modelscope.__version__}')
print(f'✓ CUDA可用: {torch.cuda.is_available()}')
if torch.cuda.is_available():
    print(f'✓ GPU数量: {torch.cuda.device_count()}')
    print(f'✓ 当前GPU: {torch.cuda.get_device_name(0)}')
"
    echo ""
}

# 创建必要的目录
create_directories() {
    print_status "创建目录结构..."
    mkdir -p ./model
    mkdir -p ./data
    mkdir -p ./output
    mkdir -p ./logs
    echo "✓ 目录创建完成"
    echo ""
}

# 下载模型
download_model() {
    local model_path="./model/unsloth/Qwen3-32B-unsloth-bnb-4bit"
    
    if [ -d "$model_path" ] && [ "$(ls -A $model_path)" ]; then
        print_warning "模型目录已存在且不为空，跳过下载"
        return 0
    fi
    
    print_status "开始下载Qwen3-32B-unsloth-bnb-4bit模型..."
    python scripts/prepare_dataset.py --dataset qwen-sft --num_samples 0 2>/dev/null || true
    
    # 使用Python脚本下载模型
    python -c "
import os
from modelscope import snapshot_download

print('开始下载模型...')
try:
    model_path = snapshot_download(
        model_id='unsloth/Qwen3-32B-unsloth-bnb-4bit',
        cache_dir='./model',
        revision='master'
    )
    print(f'✓ 模型下载完成: {model_path}')
except Exception as e:
    print(f'✗ 模型下载失败: {e}')
    print('请检查网络连接或手动下载模型')
    exit(1)
"
    echo ""
}

# 下载数据集
download_dataset() {
    local data_path="./data/chinese-qwen-sft"
    
    if [ -d "$data_path" ] && [ "$(ls -A $data_path)" ]; then
        print_warning "数据集已存在，跳过下载"
        return 0
    fi
    
    print_status "下载训练数据集..."
    python scripts/prepare_dataset.py --dataset qwen-sft --num_samples 5000
    echo ""
}

# 运行微调
run_training() {
    print_status "开始模型微调..."
    
    # 创建日志文件
    local log_file="./logs/training_$(date +%Y%m%d_%H%M%S).log"
    
    print_status "训练日志将保存到: $log_file"
    
    # 运行训练，同时输出到控制台和日志文件
    python scripts/finetune_qwen3.py \
        --model_path ./model/unsloth/Qwen3-32B-unsloth-bnb-4bit \
        --data_path ./data/chinese-qwen-sft \
        --output_dir ./output \
        --max_steps 100 \
        --batch_size 1 \
        --learning_rate 2e-4 \
        --max_seq_length 2048 \
        2>&1 | tee "$log_file"
    
    if [ ${PIPESTATUS[0]} -eq 0 ]; then
        print_status "✓ 微调完成！"
    else
        print_error "✗ 微调过程中出现错误，请检查日志文件: $log_file"
        exit 1
    fi
    echo ""
}

# 测试模型
test_model() {
    print_status "测试微调后的模型..."
    
    if [ ! -d "./output/lora_adapters" ]; then
        print_error "未找到微调后的模型，请先完成微调"
        exit 1
    fi
    
    python scripts/test_model.py \
        --model_path ./output/lora_adapters \
        --mode test
    
    echo ""
}

# 显示使用说明
show_usage() {
    echo "Qwen3 32B 微调快速启动脚本"
    echo ""
    echo "用法: $0 [选项]"
    echo ""
    echo "选项:"
    echo "  --skip-download-model    跳过模型下载"
    echo "  --skip-download-data     跳过数据集下载"
    echo "  --skip-training          跳过训练过程"
    echo "  --skip-test             跳过模型测试"
    echo "  --steps NUM             设置训练步数 (默认: 100)"
    echo "  --batch-size NUM        设置批次大小 (默认: 1)"
    echo "  --learning-rate RATE    设置学习率 (默认: 2e-4)"
    echo "  --help                  显示此帮助信息"
    echo ""
    echo "示例:"
    echo "  $0                                   # 运行完整流程"
    echo "  $0 --skip-download-model             # 跳过模型下载"
    echo "  $0 --steps 200 --batch-size 2       # 自定义训练参数"
}

# 主函数
main() {
    # 默认参数
    local skip_download_model=false
    local skip_download_data=false
    local skip_training=false
    local skip_test=false
    local max_steps=100
    local batch_size=1
    local learning_rate=2e-4
    
    # 解析命令行参数
    while [[ $# -gt 0 ]]; do
        case $1 in
            --skip-download-model)
                skip_download_model=true
                shift
                ;;
            --skip-download-data)
                skip_download_data=true
                shift
                ;;
            --skip-training)
                skip_training=true
                shift
                ;;
            --skip-test)
                skip_test=true
                shift
                ;;
            --steps)
                max_steps="$2"
                shift 2
                ;;
            --batch-size)
                batch_size="$2"
                shift 2
                ;;
            --learning-rate)
                learning_rate="$2"
                shift 2
                ;;
            --help)
                show_usage
                exit 0
                ;;
            *)
                print_error "未知参数: $1"
                show_usage
                exit 1
                ;;
        esac
    done
    
    echo "========================================"
    echo "   Qwen3 32B 微调快速启动脚本"
    echo "========================================"
    echo ""
    
    # 执行检查
    check_cuda
    check_python
    create_directories
    
    # 执行主要流程
    if [ "$skip_download_model" = false ]; then
        download_model
    else
        print_warning "跳过模型下载"
    fi
    
    if [ "$skip_download_data" = false ]; then
        download_dataset
    else
        print_warning "跳过数据集下载"
    fi
    
    if [ "$skip_training" = false ]; then
        # 更新训练参数
        sed -i "s/--max_steps [0-9]*/--max_steps $max_steps/" scripts/finetune_qwen3.py || true
        sed -i "s/--batch_size [0-9]*/--batch_size $batch_size/" scripts/finetune_qwen3.py || true
        sed -i "s/--learning_rate [0-9e.-]*/--learning_rate $learning_rate/" scripts/finetune_qwen3.py || true
        
        run_training
    else
        print_warning "跳过训练过程"
    fi
    
    if [ "$skip_test" = false ]; then
        test_model
    else
        print_warning "跳过模型测试"
    fi
    
    echo "========================================"
    print_status "所有步骤完成！"
    echo ""
    echo "微调后的模型保存在: ./output/lora_adapters"
    echo "训练日志保存在: ./logs/"
    echo ""
    echo "下一步操作："
    echo "1. 测试模型: python scripts/test_model.py --mode chat"
    echo "2. 部署模型: 参考文档中的部署章节"
    echo "========================================"
}

# 运行主函数
main "$@"
